from typing import Any, Dict, List, Union

from ..shard.placement_types import (
    DeviceMesh,Replicate,PlacementStrategy,Shard,DTensorSpec,Shard,ShardingStrategy,Replicate,Shard
)
from ..shard.embedding_strategy import embedding_strategy
import copy
from .strategy_generator import StrategyGenerator


class EmbeddingStrategyGenerator(StrategyGenerator):
    """
    EmbeddingStrategyGenerator is a generic class to generate strategies for nn.Embedding or F.embedding.
    The operation data is defined as `output = input x other`.
    """

    # def validate(self) -> bool:
    #     return super().validate()

    def collate_strategies(self) -> List[ShardingStrategy]:
        strategy_list = []
        # 只有复制 TODO
        weight_shape = self.op_data['weight'].data.shape
        input_shape = self.op_data['input'].data.shape
        output_shape = self.op_data['output'].data.shape
        # R,R->R
        # S[0],R->R
        
        weight_spec = DTensorSpec(placements =(Replicate(),),tensor_meta=self.op_data['weight'].data,mesh=self.device_mesh)
        input_spec = DTensorSpec(placements =(Replicate(),),tensor_meta=self.op_data['input'].data,mesh=self.device_mesh)
        output_spec = DTensorSpec(placements =(Replicate(),),tensor_meta=self.op_data['output'].data,mesh=self.device_mesh)
        input_specs = [weight_spec,input_spec]
        strategy = PlacementStrategy(input_specs=input_specs,output_specs=output_spec)
        strategy_list.append(ShardingStrategy(name=str(strategy),sharding_specs=strategy,compute_cost=0))
        input_specs[0].placements=(Shard(0),)
        strategy_list.append(ShardingStrategy(name=str(strategy),sharding_specs=strategy,compute_cost=0))
        return strategy_list
        # mm_strategy = embedding_strategy(self.device_mesh,weight_shape,input_shape,output_shape)
        # for strategy in mm_strategy.strategies:
        #     strategy.input_specs[0].tensor_meta = self.op_data['weight'].data
        #     strategy.input_specs[1].tensor_meta = self.op_data['input'].data
        #     strategy.output_specs.tensor_meta = self.op_data['output'].data
        #     sharding_specs = PlacementStrategy(input_specs=strategy.input_specs,output_specs=strategy.output_specs)
        #     strategy_list.append(ShardingStrategy(name=str(sharding_specs),sharding_specs=sharding_specs,compute_cost=0))
        # return strategy_list


